
#pragma once

#include <assert.h>
#include <atomic>
#include <mutex>
#include <set>
#include <memory>
#include <functional>
#include "db/DBConfig.h"


template<typename DB>
class DBPool
{
    typedef DBPool<DB> _MyType;
public:
    typedef std::shared_ptr<DB> ConnPtr;

    DBPool():init_(false){}

    ~DBPool()
    {
        for (auto conn : active_conn_){
            delete conn;
        }
        for (auto conn : free_conn_) {
            delete conn;
        }
    }

    bool init(const DBConfig& conf, size_t conn_num, size_t max_conn_num) 
    {
        assert(init_ == false);
        assert(conn_num <= max_conn_num);
        max_conn_num_ = max_conn_num;
        config_ = conf;
        bool bret = true;
        for (size_t i = 0; i < conn_num; ++i)
        {
            DB* conn = new DB;
            if (!conn->open(conf)) {
                bret = false;
            }
            free_conn_.insert(conn);
        }
        init_ = true;
        return bret;
    }

    ConnPtr get_conn()
    {
        assert(init_ = true);
        std::unique_lock<std::mutex> lck(free_mutex_);
        if (!free_conn_.empty()) {
            auto conn = *free_conn_.begin();
            free_conn_.erase(conn);
            lck.unlock();
            std::unique_lock<std::mutex> _(active_mutex_);
            active_conn_.insert(conn);
            return ConnPtr(conn, std::bind(&_MyType::_put_back, this, std::placeholders::_1));
        }
        
        std::unique_lock<std::mutex> _(active_mutex_);
        if (active_conn_.size() + free_conn_.size() >= max_conn_num_) {
            printf("active_conn_.size():%d, free_conn_.size():%d, max_conn_num_:%d\n",
                active_conn_.size(), free_conn_.size(), max_conn_num_);
            return nullptr;
        }
        lck.unlock();
        DB* conn = new DB;
        conn->open(config_);
        active_conn_.insert(conn);

        return ConnPtr(conn, std::bind(&_MyType::_put_back, this, std::placeholders::_1));
    }

private:
    void _put_back(DB* conn)
    {
        std::unique_lock<std::mutex> lck(free_mutex_);
        free_conn_.insert(conn);
        lck.unlock();
        std::unique_lock<std::mutex> _(active_mutex_);
        active_conn_.erase(conn);
    }
private:
    std::atomic<bool>           init_;
    std::mutex                  active_mutex_;
    std::set<DB*>               active_conn_;
    std::mutex                  free_mutex_;
    std::set<DB*>               free_conn_;
    DBConfig                    config_;
    size_t                      max_conn_num_;
};